import torch
import torch.nn as nn
from AlexNetModel import AlexNet
from VggNetModel import Vgg16Net
from ResNetModel import ResNet18, ResBlock


def load_alexnet(num_classes=10, in_channels=3):
    model = AlexNet(num_class=num_classes, channels=in_channels)
    return model

def load_vggnet(num_classes=10, in_channels=3):
    ratio = 8
    small_conv_struct = [(1, in_channels, 64//ratio), (1, 64//ratio, 128//ratio),
                    (2, 128//ratio, 256//ratio), (2, 256//ratio, 512//ratio), (2, 512//ratio, 512//ratio)]
    model = Vgg16Net(conv_struct=small_conv_struct, fc_features=7 * 7 * 512 // ratio, fc_hidden_num=4096 // ratio, num_class=num_classes, channels=in_channels)
    return model

def load_resnet(num_classes=10, in_channels=3):
    model = ResNet18(ResBlock, num_class=num_classes, channels=in_channels)
    return model